"""力预测"""
import torch
from torch import nn
import torch.nn.functional as F
from cganet import GCNF
import numpy as np
from math import sqrt

#parameter
num_epochs = 200
batch_size = 64
learning_rate = 0.001
device_ids = [0,1]
torch.cuda.set_device(device_ids[0])

#固定随机数种子，使得结果可复现
"""np.random.seed(0)
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
torch.backends.cudnn.deterministic = True"""

#define
torch.set_default_dtype(torch.float32)
device = torch.device('cuda:1' if torch.cuda.is_available() else 'cpu')
model = GCNF(32,24,18,26,30,32)
model = nn.DataParallel(module=model, device_ids=device_ids)
model = model.cuda(device_ids[0])

criterion = nn.L1Loss()
#test_criterion = nn.MSELoss(reduction='mean')
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

#load data
data = np.load('/home/hp/users/fuli/11_Ag_large/train/Ag34_CDEFGH3.npz')
temp = data['D']
spl = int(temp.shape[0]*0.9) #训练集和测试集划分
D = temp[:spl] #距离矩阵
D_test = temp[spl:]
C = data['C'][:spl] #坐标矩阵
C_test = data['C'][spl:]
E = data['E'][:spl] #能量
E_test = data['E'][spl:]
F = data['F'][:spl] #力
F_test = data['F'][spl:]
print(C.shape, C_test.shape)

#train
task_name = 'Ag34f_32,24,18,26,30,32'
loss_file = open(task_name+'.txt','w')
loss_file.write(open('force_train.py',encoding='utf-8').read()+'\n')
last_best_test = 1e10

for epoch in range(num_epochs):
    total_loss = 0
    for i in range(D.shape[0]//batch_size):
        c = torch.tensor(C[i*batch_size:(i+1)*batch_size]).cuda(device_ids[0])
        d = torch.tensor(D[i*batch_size:(i+1)*batch_size]).cuda(device_ids[0])
        f = torch.tensor(F[i*batch_size:(i+1)*batch_size]).cuda(device_ids[0])
        out_f = model(c,d)
        loss = criterion(f, out_f)
        #loss = criterion(f[...,0], out_f[...,0]) + criterion(f[...,1], out_f[...,1]) + criterion(f[...,2], out_f[...,2])
        total_loss += loss.item()
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        if i%500==0:
            print ('    Batch [{}/{}], force loss={:.6f}'.format(
                    i+1, D.shape[0]//batch_size, total_loss/(i+1)))
    
    info = 'Epoch [{:>3d}/{}], mean force loss={:.8f}'.format(
            epoch+1, num_epochs, total_loss/(D.shape[0]//batch_size))
    print(info)
    loss_file.write(info)
    
    #test
    total_loss = 0
    for i in range(D_test.shape[0]//batch_size):
        c = torch.tensor(C_test[i*batch_size:(i+1)*batch_size]).cuda(device_ids[0])
        d = torch.tensor(D_test[i*batch_size:(i+1)*batch_size]).cuda(device_ids[0])
        f = torch.tensor(F_test[i*batch_size:(i+1)*batch_size]).cuda(device_ids[0])
        out_f = model(c,d)
        loss = criterion(f, out_f)
        total_loss += loss.item()
    info = 'test loss={:.8f}'.format(
            total_loss/(D_test.shape[0]//batch_size))
    print(info)
   
    #save
    loss_file.write('    '+info+'\n')
    loss_file.flush()
    if total_loss < last_best_test:
        last_best_test = total_loss
        torch.save(model, task_name+'.pth')
